# main.py
import argparse, json
import datetime
import os
import logging
import torch, random

from server import *
from client import *
import models, datasets


if __name__ == '__main__':

	parser = argparse.ArgumentParser(description='Federated Learning')
	parser.add_argument('-c', '--conf', dest='conf')
	args = parser.parse_args()

	# python main.py -c ./utils/conf.json
	# 在终端输入，表示python解释器运行main.py  ///  -c ./utils/conf.json将目录下的文件赋到参数args.conf
	with open(args.conf, 'r') as f:
		conf = json.load(f)	

	train_datasets, eval_datasets = datasets.get_dataset("./data/", conf["type"])
	# 此处直接调用了datasets.py的函数得到了数据集
	
	server = Server(conf, eval_datasets)  # server是一个继承了类的服务端
	clients = []
	
	for c in range(conf["no_models"]):  # 在clients列表中依次放入十个客户端所要训练的的 数据子集
		clients.append(Client(conf, server.global_model, train_datasets, c))
		# server.global_model是全局模型，也是全局初始模型
		# 10个客户端但不同数据集的类已经通过构造函数完成

	print("\n\n")
	for e in range(conf["global_epochs"]):  # 全局训练轮次global_epochs为20

		candidates = random.sample(clients, conf["k"])  # 随机挑选k=5个客户端 参与这轮次训练全局模型

		weight_accumulator = {}  # 先提前定义本轮次的 权重差积累（wqf名词）

		for name, params in server.global_model.state_dict().items():
			weight_accumulator[name] = torch.zeros_like(params)  # 权重差积累置零

		for c in candidates:  # 5次循环，每个客户机本地训练。（此时c就代替了Client）
			diff = c.local_train(server.global_model)
			# 得到本轮次的 本客户端通过模型训练（不同客户端数据集不同）的 权重网络的变化，即权重差

			for name, params in server.global_model.state_dict().items():
				weight_accumulator[name].add_(diff[name])
		# 5次循环之后，得到五个客户端更新的模型和本轮次的全局模型之差。接下来就可以模型聚合！

		server.model_aggregate(weight_accumulator)  # 利用联邦平均算法进行模型聚合，更新全局模型（函数定义于server.py）

		acc, loss = server.model_eval()  # 模型评估

		print("Global Epoch %d, accuracy: %f, loss: %f\n" % (e+1, acc, loss))


# 1、CPU全部成功运行，花费了5个小时
# 2、GPU全部成功运行，花费了半个小时
# 3、在Pytorch，对于数据转换这一块需要特别仔细，只要某个训练中需要用到的变量（训练、评估数据和模型数据）仍然存放在CPU上，
# Pytorch就会报错，说至少发现数据存放在cpu和cuda:0两个不同的环境下。
# README.md 的编写是markdown格式
